##############################################################################
## plot.cna.bacom
##
## Description:
##   'plot.cna.bacom' plots copy number data.
## 
## Usage:
##   plot.cna.bacom(copyNumber, chrID, segMean, alpha)
##
## Arguments:
##   copyNumber: a vector of copy number profile.
##   chrID: a character vector of chromosome ID.
##   segMean: segment means.
##   alpha: percentage of normal cell contamination.
##   position: a vector of x-axis indices (consecutive numbers or physical
##             locations).
##   ylim: y-axis limit.
##   chrToPlot: a character vector of chromosome IDs to be plotted.
##   isPlotSegMeanBeforeCorrection: TRUE if to plot segment mean before 
##                                  correction.
##   isShowLegend: TRUE if to show legend on the figure.
##   isShowAlpha: TRUE if to show the value of alpha on the figure.
##
## Example:
##   plot.cna.bacom(copyNumber, chrID, segMean, alpha)
##
## Author:
##   Bai Zhang    baizhang@vt.edu
##
##############################################################################


plot.cna.bacom <- function(copyNumber, chrID, segMean, alpha, position = 1:length(copyNumber), ylim = c(-1, 6),
                           chrToPlot=paste(1:22), isPlotSegMeanBeforeCorrection = TRUE, isShowLegend = TRUE,
                           isShowAlpha = TRUE) {

  for (chr in intersect(unique(chrID), chrToPlot)) {

    cn <- copyNumber[chrID == chr]
    segMean.chr <- segMean[chrID == chr]

    position.chr <- position[chrID == chr]
    if (sum(as.numeric(position) - as.numeric(1:length(copyNumber))) == 0) {
      position.chr <- position.chr - position.chr[1] + 1
    }

    plot(x=position.chr, y=cn, pch = ".", col = "#999999", ylim = ylim, xlab = paste("Chr ", chr), ylab = "Copy number", cex.lab=1.2)

    if (isShowAlpha) {
      text(x = ((max(position.chr) - min(position.chr))*0.9+min(position.chr)), y = max(ylim)*0.9,
           labels = bquote(bar(alpha) == .(round(alpha, 3))), cex = 1.8)
    }
    
    
    grid()
    
    segment <- profile2segment(segMean.chr)

    
    for (i in 1:length(segment$seg.mean)) {
      if (segment$seg.mean[i] < 1.9) {
        seg.color <- "#33a3ff"
        seg.mean.color <- "#0000FF"
      } else if (segment$seg.mean[i] > 2.1) {
        seg.color <- "#bd5adf"
        seg.mean.color <- "#FF0000"
      } else {
        seg.color <- "#999999"
        seg.mean.color <- "#00FF00"
      }



      if (segment$seg.mean[i] < 1.9 || segment$seg.mean[i] > 2.1) {
        corrected.segMean <- (segment$seg.mean[i] - 2*alpha) / (1 - alpha)
      } else {
        corrected.segMean <- 2
      }

      s <- segment$seg.start[i]
      e <- segment$seg.end[i]
      points(x=position.chr[s:e], y = cn[s:e], pch=".", col = seg.color)

      if (isPlotSegMeanBeforeCorrection) {
        lines(x=position.chr[c(segment$seg.start[i], segment$seg.end[i])], y=c(segment$seg.mean[i], segment$seg.mean[i]), col=seg.mean.color, lwd=2, lty = 3)
      }
      lines(x=position.chr[c(segment$seg.start[i], segment$seg.end[i])], y=c(corrected.segMean, corrected.segMean), col=seg.mean.color, lwd=2, lty = 1)


      

    }

    if (isShowLegend) {
      if (isPlotSegMeanBeforeCorrection) {
        legend(0, max(ylim), c("Segment mean: ", "deletion after correction", "deletion before correction", "amplification after correction", "amplification before correction", "normal"), col = c("white", "blue", "blue", "red", "red", "green"), lty = c(0, 1, 3, 1, 3, 1), lwd = c(1, 2, 2, 2, 2, 2), cex = 0.8)
      } else {
        legend(0, max(ylim), c("Segment mean: ", "deletion after correction", "amplification after correction", "normal"), col = c("white", "blue", "red", "green"), lty = c(0, 1, 1, 1), lwd = c(1, 2, 2, 2), cex = 0.8)
      }
    }

    
    
  }
  
  
}


profile2segment <- function(segMean) {

  x <- segMean
  x.shifted <- c(x[-1], x[length(x)])
  x <- x - x.shifted
  seg.start <- c(1, which(x != 0)+1)
  seg.end <- c(which(x != 0), length(x))
  seg.mean <- c(segMean[seg.start])

  segment <- list(seg.start=seg.start, seg.end=seg.end, seg.mean=seg.mean)

  segment
  
}

